import pandas as pd
import pyodbc
import datetime
import warnings
import os
import time
import gc

# 忽略pandas警告
warnings.filterwarnings('ignore', category=UserWarning)

def get_current_shift_info():
        """获取当前的班次信息（日期、班次）"""
        now = datetime.datetime.now()
        hour = now.hour + now.minute / 60
        today = datetime.datetime.today()
        
        if hour < 8.5:
            # 早班之前，属于前一天的夜班
            today = today - datetime.timedelta(days=1)
            dateiv = today.strftime("%Y-%#m-%#d")
            datetpe = today.strftime("%Y-%m-%d")
            datef=today.strftime("%Y%m%d")
            shift = "N"
        elif hour > 20.5:
            # 晚班时间
            dateiv = today.strftime("%Y-%#m-%#d")
            datetpe = today.strftime("%Y-%m-%d")
            datef=today.strftime("%Y%m%d")
            shift = "N"
        else:
            # 白班时间
            dateiv = today.strftime("%Y-%#m-%#d")
            datetpe = today.strftime("%Y-%m-%d")
            datef=today.strftime("%Y%m%d")
            shift = "D"

        return dateiv,datetpe, shift,datef
# 尝试使用GBK编码读取CSV文件，这是中文Windows系统常见的编码

def read_mdb(access_path):
    """读取MDB文件并返回pandas DataFrame"""
    conn_str = (
        r"Driver={Microsoft Access Driver (*.mdb, *.accdb)};"
        rf"Dbq={access_path};"
    )
    conn = pyodbc.connect(conn_str)
    
    try:
        # 使用简化的SQL查询，避免特殊字符和参数问题
        # df = pd.read_sql("SELECT UniqueId,TestTime,CustomerLine,Class,Bin,Isc,Uoc,FF,Eta,Insol,Tcell,Tmonicell,RshuntDfDr,RserLfDfIEC,Irev2,IvGrade,Aoi1Q,Aoi1R,Aoi2Q,Aoi2R,BinFileName,Busbar,CalibrationValue,CellTempOffset FROM Results", conn)
        df = pd.read_sql("SELECT UniqueId,TestTime,CustomerLine,Class,Bin,Isc,Uoc,FF,Eta,Insol,Tcell,Tmonicell,Rshunt,Rser,Irev2,IvGrade,Aoi1Q,Aoi1R,Aoi2Q,Aoi2R,BinFileName,Busbar,CalibrationValue,CellTempOffset FROM Results", conn)
        
        return df
    except Exception as e:
        raise e
    finally:
        # 确保连接关闭
        conn.close()

def optimize_for_large_data(df):
    """优化数据类型以减少内存使用"""
    # 针对bin列进行专门优化
    if 'bin' in df.columns:
        # 尝试将bin列转换为更紧凑的类型
        try:
            # 检查是否可以转换为整数
            df['bin'] = pd.to_numeric(df['bin'], downcast='integer')
        except ValueError:
            # 如果失败，尝试转换为分类类型（如果是字符串）
            if df['bin'].dtype == 'object':
                df['bin'] = df['bin'].astype('category')
    return df

def main(use_compression=True):
    """主函数
    
    Args:
        use_compression: 是否使用数据压缩，默认为True
    """
    print("===== 开始数据处理任务 =====")
    start_time = time.time()
    print(f"✓ 压缩设置: {'启用' if use_compression else '禁用'}")
    
    # 读取IP配置
    try:
        ip = pd.read_csv('IP.csv', encoding='gbk')
        print(f"✓ 成功读取IP配置，共 {len(ip)} 个数据源")
    except Exception as e:
        print(f"❌ 读取IP配置失败: {e}")
        return
    
    # 获取当前班次信息
    dateiv, datetpe, shift, datef = get_current_shift_info()
    print(f"✓ 当前班次信息：日期={dateiv}, 班次={shift}, 表名={datef}")
    
    # 初始化变量
    total_sources = len(ip)
    success_count = 0
    total_rows = 0
    table_name = datef
    parquet_file = f'{table_name}.parquet'
    
    # 先检查文件是否存在，如果存在先删除
    if os.path.exists(parquet_file):
        os.remove(parquet_file)
        print(f"✓ 删除已存在的文件: {parquet_file}")
    
    # 单批处理，直接写入文件，避免累积所有数据
    process_start_time = time.time()
    first_write = True  # 标记是否是第一次写入
    
    # 减小批次大小以减少内存占用
    for i in range(total_sources):
    # for i in range(0,1):
        try:
            ip_address = ip.iloc[i]['IP']
            ip_path = ip.iloc[i]['Path']
            ip_line = ip.iloc[i]['Line']
            access_path = f"{ip_address}{ip_path}\\{dateiv}-{ip_line}-D.mdb"
            
            print(f"[{i+1}/{total_sources}] 正在读取: {access_path}")
            
            # 读取MDB文件
            df = read_mdb(access_path)
            current_rows = len(df)
            total_rows += current_rows
            
            # 为每条记录标记来源
            df['source_ip'] = ip_address
            df['source_line'] = ip_line
            df['source_file'] = access_path
            
            # 优化数据类型
            df = optimize_for_large_data(df)
            
            # 直接写入文件，避免累积数据
            write_mode = 'w' if first_write else 'a'
            
            # 优化的Parquet写入参数，启用压缩以减小文件体积
            # 对于不支持append参数的pandas版本，使用替代方案
            if first_write:
                # 首次写入，直接创建文件
                df.to_parquet(
                    parquet_file,
                    engine='pyarrow',
                    index=False,
                    compression='snappy' if use_compression else 'none',  # 根据参数决定是否压缩
                    row_group_size=1000000,
                    version='2.4',
                    use_dictionary=False
                )
            else:
                # 后续写入，使用临时文件合并方式
                temp_file = f'{parquet_file}.temp'
                # 先写入临时文件
                df.to_parquet(
                    temp_file,
                    engine='pyarrow',
                    index=False,
                    compression='snappy' if use_compression else 'none',
                    row_group_size=1000000,
                    version='2.4',
                    use_dictionary=False
                )
                # 读取现有文件
                existing_df = pd.read_parquet(parquet_file, engine='pyarrow')
                # 合并数据
                combined_df = pd.concat([existing_df, df], ignore_index=True)
                # 写回原文件
                combined_df.to_parquet(
                    parquet_file,
                    engine='pyarrow',
                    index=False,
                    compression='snappy' if use_compression else 'none',
                    row_group_size=1000000,
                    version='2.4',
                    use_dictionary=False
                )
                # 删除临时文件
                if os.path.exists(temp_file):
                    os.remove(temp_file)
                # 清理内存
                del existing_df, combined_df
            
            success_count += 1
            print(f"  ✓ 读取并写入成功，共 {current_rows} 行数据")
            
            # 更新写入状态
            if first_write:
                first_write = False
            
            # 立即清理内存
            del df
            gc.collect()  # 强制垃圾回收
            
        except Exception as e:
            print(f"  ❌ 处理失败: {e}")
            # 继续处理下一个文件
            continue
    
    process_end_time = time.time()
    print(f"✓ 数据读取和写入完成，耗时: {process_end_time - process_start_time:.2f} 秒")
    print(f"✓ 数据源处理完成：成功 {success_count}/{total_sources}")
    print(f"✓ 总写入数据行数: {total_rows}")
    
    # 计算并显示总耗时
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"✓ 任务总耗时: {elapsed_time:.2f} 秒 ({elapsed_time/60:.2f} 分钟)")
    print("===== 数据处理任务结束 =====")

# 确保在Windows下正确处理多进程（如果需要扩展）
if __name__ == "__main__":
    # 可以根据需要设置是否使用压缩
    # 压缩会减小文件体积但可能稍微增加处理时间
    # main(use_compression=False)  # 禁用压缩以提高速度
    main(use_compression=True)  # 启用压缩以减小文件体积
